# based on SRE_inv functions

# Initialises a matrix with n individuals
# individuals have a location (spX, spY) and are 
# initilised with gene for p.disp
# based on Stable range edge functions
init.inds<-function(n, spX, spY, hvar=TRUE){
	X<-runif(n, 1, spX+1)%/%1 #grid ref
	Y<-runif(n, 1, spY+1)%/%1
	if (hvar==FALSE) H<-rep(1, n)
	else H<-runif(n, 0, 1) #Habitat match
	cbind(X, Y, H)		
}

# initialises a habitat matrix based on homogenous area (length=height=core) and nsteps
# nsteps defines the rate at which habitat declines beyond homogenous core area and >0
# homogenous area has default habitat value of 1
# gradient ends five steps from the edges of the space
init.hab<-function(core, nsteps){
	if (nsteps==0) return(matrix(data=1, nrow=core, ncol=core))
	ncol1<-5
	simple<-matrix(0, nrow=core, ncol=ncol1)
	simple2<-matrix(1, nrow=core, ncol=core)
	d<-seq(1, 0, length.out=nsteps+1)[-1]
	d<-rep(d, times=core)
	grade1<-t(matrix(d, ncol=core, nrow=nsteps))
	d<-rev(seq(1, 0, length.out=nsteps+1)[-1])
	d<-rep(d, times=core)
	grade2<-t(matrix(d, ncol=core, nrow=nsteps))
	t(cbind(simple, grade2, simple2, grade1, simple)) #same format as density matrix
}

# hassel commins population growth where b=1 (contest competition)
# generates expected fecundity
hass.comm<-function(lambda, K, N){
	exp<-lambda/(1+(lambda-1)*N/K)
}

# reproduces individuals based on density and habitat sepcific fitness, kills parents
# kills dispersing offspring based on dispersal cost
# disperses individuals based on fixed dprob
repro.disp<-function(popmatrix, hab.matrix, spX, spY, K, lambda, d.cost, d.prob, n.list, mutn, adap=TRUE, impact=0, k){
	density<-table(factor(popmatrix[,"X"], levels=1:spX), factor(popmatrix[,"Y"], levels=1:spY)) # adult density through space
	temp<-(popmatrix[,"Y"]-1)*spX+popmatrix[,"X"] #matrix indexes for density/habitat matrices
	dens<-density[temp] # densities returned to each individual
	habval<-hab.matrix[temp] # habitat returned to each individual
	if (adap==TRUE) hab<-1-abs(habval-popmatrix[,"H"]) #decline in survival with habitat fit
	hab<-ifelse(hab<0, 0, hab)
	offs<-rpois(length(popmatrix[,1]), hab*hass.comm(lambda, K, dens)) # weaned offspring based on fitness and density
	if (sum(offs)==0) return(NULL) #catch extinctions
	inds<-1:length(offs) #vector indexes for offspring
	inds<-rep(inds, times=offs)
	popmatrix<-popmatrix[inds,1:3] #replace parents with offspring. Offspring inherit location and P from parents
	if (class(popmatrix)=="numeric") popmatrix<-matrix(popmatrix, ncol=3, dimnames=list(NULL, c("X", "Y", "H"))) #cases where only one individual left
	if (impact>0){
		temp<-(popmatrix[,"Y"]-1)*spX+popmatrix[,"X"] #matrix indexes for density/habitat matrices
		habval<-hab.matrix[temp] # habitat returned to each individual
		imp<-ifelse(imp.fun(habval, impact, k)>1, 1, imp.fun(habval, impact, k))
		imp<-rbinom(length(habval), 1, imp) # random mortality associated with impact
		popmatrix<-subset(popmatrix, imp==0)
		if (length(popmatrix)==0) return(NULL) #catch extinctions (again)
	}
	mut.sample<-rbinom(nrow(popmatrix), 1, mutn) # individuals copping a mutation
	inds<-which(mut.sample==1)
	if (length(inds)>0){
		mutation<-rnorm(length(inds), popmatrix[inds, "H"], sd=0.05)
		mutation<-ifelse(mutation>1, 1, mutation)
		mutation<-ifelse(mutation<0, 0, mutation)
		popmatrix[inds, "H"]<-mutation
		
	}
	disp<-rbinom(nrow(popmatrix), 1, d.prob) #disperse or not
	popmatrix<-cbind(popmatrix, disp)
	mort<-d.cost*disp #probability of dying during dispersal
	popmatrix<-subset(popmatrix, rbinom(length(mort), 1, mort)!=1) # remove those dying during dispersal
	popmatrix<-disperse(popmatrix, n.list, spX)
	popmatrix
}

#calculates neighbours on a grid excluding nonsense neighbours and returns a list of matrices of grid refs
#doesn't include source cell
# space is cylindrical
neighbours.init<-function(spX, spY, torus=TRUE){
	neigh<-vector("list", length=spX*spY)
	if (torus==F){
		for (x in 1:spX){
			for (y in 1:spY){
				out<-matrix(nrow=8, ncol=2)
				ifelse((x-1)==0, X<-NA, X<-x-1) #correct for indices going to zero
				ifelse((y-1)==0, Y<-spY, Y<-y-1)
				ifelse((y+1)>spY, Yp<-1, Yp<-y+1) # correct for indices going to > space
				ifelse((x+1)>spX, Xp<-NA, Xp<-x+1)
				# Assign neighbours
				out[1,]<-c(Xp, y)
				out[2,]<-c(X, y)
				out[3,]<-c(x, Y)
				out[4,]<-c(Xp, Y)	
				out[5,]<-c(X, Y)
				out[6,]<-c(x, Yp)
				out[7,]<-c(Xp, Yp)
				out[8,]<-c(X, Yp)
				out<-subset(out, !is.na(apply(out, 1, sum)))
				neigh[[(y-1)*spX+x]]<-out
			}	
		}
	}
	else{
		for (x in 1:spX){
			for (y in 1:spY){
				out<-matrix(nrow=8, ncol=2)
				ifelse((x-1)==0, X<-spX, X<-x-1) #wrap both axes
				ifelse((y-1)==0, Y<-spY, Y<-y-1)
				ifelse((y+1)>spY, Yp<-1, Yp<-y+1)
				ifelse((x+1)>spX, Xp<-1, Xp<-x+1)
				# Assign neighbours
				out[1,]<-c(Xp, y)
				out[2,]<-c(X, y)
				out[3,]<-c(x, Y)
				out[4,]<-c(Xp, Y)	
				out[5,]<-c(X, Y)
				out[6,]<-c(x, Yp)
				out[7,]<-c(Xp, Yp)
				out[8,]<-c(X, Yp)
				out<-subset(out, !is.na(apply(out, 1, sum)))
				neigh[[(y-1)*spX+x]]<-out
			}	
		}	
	}
	
	neigh	#return the list
}

# samples a matrix and returns one row of the matrix
matrix.sample<-function(mtrx, size=1){
	mtrx[sample(seq(length(mtrx[,1])), size=size),]	
}

# executes nearest neighbour dispersal where individuals disperse to neighbouring cell with probability=prob.d 
# is called after reproduction/survival
disperse<-function(popmatrix, n.list, spX){
	if (length(popmatrix[,3])==0) return(popmatrix)
	disp<-which(popmatrix[,"disp"]==1) #get indices of surviving dispersers
	if (length(disp)==0) return(popmatrix)
	sub.list<-n.list[(popmatrix[disp,"Y"]-1)*spX+popmatrix[disp,"X"]] #generate list of neighbour matrices
	sub.list<-lapply(sub.list, matrix.sample) #samples one row from each neighbour matrix
	sub.list<-do.call("rbind", sub.list) #places result into a matrix
	popmatrix[disp, 1:2]<-sub.list #assigns new locations back to popmatrix
	popmatrix
}

# plots population density and mean trait values through space
plotter<-function(popmatrix, spX, spY, K){
	density<-table(factor(popmatrix[,"X"], levels=1:spX), factor(popmatrix[,"Y"], levels=1:spY)) # adult density through space
	density<-density/K
	hval<-tapply(popmatrix[,"H"], list(factor(popmatrix[,"X"], levels=1:spX), factor(popmatrix[,"Y"], levels=1:spY)), mean) #mean trait values through space
	par(mfrow=c(2,1))
	image(density, zlim=c(0,2), col=(grey((12:0)/12)), main="Population density")
	image(hval, zlim=c(0,1.1), col=(grey((12:0)/12)), main="Habitat values")
}

# samples a column and returns mean and variance of trait values
col.sample<-function(pop, column){
	rws<-which(pop[,"X"]==column) 
	mean.h<-mean(pop[rws, "H"])
	var.h<-var(pop[rws,"H"])	
	out<-c(mean.h, var.h)
	out
}

imp.fun<-function(habval, impval=1, k){
	impval*exp(-k*(1-habval))	
}

imp.fun2<-function(habval, impval=1, k){
	impval*(1-exp(-k*habval))	
}

#mother function
mother<-function(n=20*K*10, ngens.init=100, gens.breach=3000, spY=20, steps=10, K=40, lambda, d.cost, d.prob, mutn=0.05, plot=F, im=0.5, k=5){
	pop<-init.inds(n, spY, spY, hvar=T)
	maxX<-10+2*steps+spY
	hab<-init.hab(spY, 0)	
	nbrs<-neighbours.init(spY, spY) #reveal all neighbours
	for (i in 1:ngens.init){
		pop<-repro.disp(pop, hab, spY, spY, K, lambda, d.cost, d.prob, nbrs, mutn)
		if (plot==T) {
			temp=pop
			temp[,"X"]=temp[,"X"]+5+steps
			plotter(temp, maxX, spY, K)	
		}
	}
	pop[,"X"]=pop[,"X"]+steps+5
	hab<-init.hab(spY, steps)		# re-assign habitat
	nbrs<-neighbours.init(maxX, spY) #reveal all neighbours
	ngens<-ngens.init+gens.breach
	st.time<-i
	for (i in (ngens.init+1):(ngens.init+11)){
		pop<-repro.disp(pop, hab, maxX, spY, K, lambda, d.cost, d.prob, nbrs, mutn, impact=0)
		if (plot==T) plotter(pop, maxX, spY, K)
	}
	kill<-sample(1:maxX^2, size=maxX^2*im)
	for (i in (ngens.init+11):ngens){
		temp<-(pop[,"Y"]-1)*maxX+pop[,"X"]
		pop<-subset(pop, !temp%in%kill)
		pop<-repro.disp(pop, hab, maxX, spY, K, lambda, d.cost, d.prob, nbrs, mutn, impact=0)
		if (plot==T) plotter(pop, maxX, spY, K)
		if (max(pop[,"X"])==maxX | i==ngens | min(pop[,"X"])==1) {
			b.time<-i
			out<-c(st.time, b.time)
			names(out)<-c("start", "breach")
			return(out)
		}	
	}
	
}

